/*
 * SimRewardFunction.cpp
 */

#include "SimRewardFunction.h"

SimRewardFunction::SimRewardFunction(SimEnvironment *psimEnv, int pexptNum)
{
	simEnv = psimEnv;
	exptNum = pexptNum;

}

double SimRewardFunction::getReward(CStateCollection * oldState, CAction *action, CStateCollection *newState){

	switch(exptNum){

	case EXPERIMENT_1_SPHERE_GOAL:
		return defaultRewardFast(newState);
		break;

	case EXPERIMENT_2_PUSH_BALL:
		//return defaultRewardFast(newState);
		switch(rewardType){

			case REWARD_POTENTIAL: {
				// potential reward function
				// distance of ball from goal.
				return defaultRewardPotential(oldState, newState, 2);
				break;
			}
			case REWARD_DEFAULT:
				return defaultRewardFast(newState);
				break;
		}
		break;

	case EXPERIMENT_3_JOINT_ACTION_PUSH: {

		switch(rewardType){

			case REWARD_POTENTIAL: {
				// potential reward function
				// distance of ball from goal.
				return defaultRewardPotential(oldState, newState, 4);
				break;
			}
			case REWARD_DEFAULT:
				return defaultRewardFast(newState);
				break;
		}
		break;
	}
	case EXPERIMENT_4_COOPERATION_OBSTACLE:
		return defaultRewardFast(newState);
		break;

	case EXPERIMENT_5_CAPTURE_HUMAN:
		return defaultRewardFast(newState);
		break;

	case EXPERIMENT_6_MULTI_AGENT_PUSH:

		switch(rewardType){

			case REWARD_POTENTIAL: {
				// potential reward function
				// distance of ball from goal.
				return defaultRewardPotential(oldState, newState, 4);
				break;
			}
			case REWARD_DEFAULT:
				return defaultRewardFast(newState);
				break;
		}
		break;

	case EXPERIMENT_7_JOINTS:
		// THIS IS DIFFERENT
		switch(rewardType){

					case REWARD_POTENTIAL: {
						// potential reward function
						// distance of ball from goal.
						return defaultRewardPotential(oldState, newState, 2);
						break;
					}
					case REWARD_DEFAULT:
						return defaultRewardFast(newState);
						break;
				}
				break;

		break;

	case EXPERIMENT_8_MULTI_IL_PUSH:

		switch(rewardType){

			case REWARD_POTENTIAL: {
				// potential reward function
				// distance of ball from goal.
				return defaultRewardPotential(oldState, newState, 2);
				break;
			}
			case REWARD_DEFAULT:
				return defaultRewardFast(newState);
				break;
		}
		break;

	}
}

double SimRewardFunction::defaultRewardPotential(  CStateCollection *oldState, CStateCollection *newState, int numState){
	if(simEnv->isThreadFailedState(newState->getState())){
		return -10000;
	}
	if(simEnv->isResetState(newState->getState())){
		return 10000;
	}

	double ballOldX = oldState->getState()->getContinuousState(numState);
	double ballNewX = newState->getState()->getContinuousState(numState);

	if(ballNewX == ballOldX){
		return -100;
	}

	return (ballNewX - ballOldX)*100;

}

double SimRewardFunction::defaultRewardFast( CStateCollection *newState){
	if(simEnv->isThreadFailedState(newState->getState())){
		return -10000;
	}
	if(simEnv->isResetState(newState->getState())){
		return 10000;
	}
	return -100;
}


double SimRewardFunction::defaultRewardNotFall( CStateCollection *newState){
	if(simEnv->isThreadFailedState(newState->getState())){
		return -10000;
	}
	if(simEnv->isResetState(newState->getState())){
		return 10000;
	}
	return 0;
}





